from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker
from sqlalchemy_utils import database_exists, create_database


from dp_analysis.config import MYSQL


class Mysql(object):
    def __init__(self, database, source=0):
        self.engine = create_engine(
            f"mysql+pymysql://"
            f"{MYSQL[source]['username']}:"
            f"{MYSQL[source]['password']}@"
            f"{MYSQL[source]['host']}:"
            f"{MYSQL[source]['port']}/"
            f"{database}?"
            f"charset=utf8"
        )
        if not database_exists(self.engine.url):
            create_database(self.engine.url)
        self.session = None

    def __enter__(self):
        Session = sessionmaker(bind=self.engine)
        self.session = Session()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        self.session.close()
